# Licensed under the MIT license.

"""
rStar推理工具模块

本模块提供了rStar推理系统的核心工具函数和类定义，包括：
1. 节点类型定义 - 定义MCTS树中不同类型的节点
2. 异常处理类 - 处理生成器相关异常
3. 解决方案跟踪处理 - 拼接和管理推理过程中的解决方案轨迹
4. 节点选择和评估 - 从MCTS树中找到最佳解决方案

主要功能：
- 支持多种推理节点类型（用户问题、子问题、直接答案等）
- 提供解决方案轨迹的拼接和掩码功能
- 实现最佳解决方案的选择算法
- 支持潜在分数评估和随机选择

使用场景：
- MCTS推理树的构建和管理
- 推理过程的跟踪和可视化
- 解决方案的评估和选择
"""

from enum import Enum, unique  # 枚举类型定义，unique确保枚举值唯一
import re  # 正则表达式库，用于文本模式匹配和提取
import math  # 数学函数库，提供数学运算（如ceil、log、sqrt等）
from typing import Dict, Tuple  # 类型注解，提供字典和元组类型标注
from colorama import Fore, Style  # 终端颜色输出库，用于彩色文本显示
from eval_src import Evaluator  # 评估器模块，用于答案评估和比较


@unique
class Node_Type(Enum):
    """
    MCTS推理树中的节点类型枚举
    
    定义了rStar推理系统中所有可能的节点类型：
    - USER_QUESTION: 用户原始问题节点（根节点）
    - REPHRASED_USER_QUESTION: 重新表述的用户问题节点
    - DIRECT_ANSWER: 直接答案节点（终端节点）
    - SUBQUESTION: 子问题节点（包含子问题和子答案）
    - RE_SUBANSWER: 重新回答子问题节点
    - OST_STEP: 一步思考节点（One-Step Thought）
    """
    USER_QUESTION = "USER_QUESTION"
    REPHRASED_USER_QUESTION = "REPHRASED_USER_QUESTION"
    DIRECT_ANSWER = "DIRECT_ANSWER"
    SUBQUESTION = "SUBQUESTION"
    RE_SUBANSWER = "RE_SUBANSWER"
    OST_STEP = "OST_STEP"


class GeneratorError(Exception):
    """
    生成器异常类
    
    用于处理在推理过程中生成子节点时发生的异常。
    记录异常发生的源头、输入和输出信息，便于调试和错误处理。
    
    Attributes:
        source (str): 异常发生的源头描述
        io_input (str): 导致异常的输入内容
        io_output_list (list): 导致异常的输出列表
    """
    def __init__(self, source, io_input, io_output_list) -> None:
        """
        初始化生成器异常
        
        Args:
            source (str): 异常源头描述
            io_input (str): 输入内容
            io_output_list (list): 输出列表
        """
        super().__init__()

        self.source = source
        self.io_input = io_input
        self.io_output_list = io_output_list


def split_user_question(user_question: str):
    """
    分割用户问题为上下文和核心问题
    
    将用户问题按最后一个句号分割为两部分：
    - 前半部分：问题的上下文信息
    - 后半部分：核心问题内容
    
    Args:
        user_question (str): 完整的用户问题
        
    Returns:
        tuple: (问题上下文, 核心问题) 的元组
    """
    user_question = user_question.strip().rstrip(".")  # 1. 去掉首尾空格和末尾句号
    last_period_id = user_question.rfind(".")          # 2. 找到最后一个句号的位置
    assert last_period_id < len(user_question) - 1     # 3. 确保句号不是最后一个字符
    user_question_context = user_question[: last_period_id + 1].strip()  # 4. 前半部分（包含句号）
    user_question_problem = user_question[last_period_id + 1 :].strip()  # 5. 后半部分（句号后的内容）
    return user_question_context, user_question_problem


def reach_terminal_subquestion(subquestion: str, user_question: str):
    """
    判断子问题是否为终端子问题
    
    终端子问题的判断条件：
    1. 子问题以"Now we can answer"开头，表示可以回答原问题
    2. 子问题包含原问题的核心部分，表示回到了原问题
    
    Args:
        subquestion (str): 当前子问题
        user_question (str): 原始用户问题
        
    Returns:
        bool: 如果是终端子问题返回True，否则返回False
    """
    assert subquestion is not None

    if "Now we can answer" in subquestion:
        #! remember that: when the original question is answerable, please start the subquestion with "Now we can answer the question: "
        return True

    user_question_2nd_part = split_user_question(user_question)[1]
    if user_question_2nd_part.lower() in subquestion.lower():
        return True

    return False


def reach_terminal_ost_step(ost_step: str):
    """
    判断一步思考是否为终端步骤
    
    终端一步思考的判断条件：
    - 步骤内容包含"answer is"，表示已得出最终答案
    
    Args:
        ost_step (str): 一步思考的内容
        
    Returns:
        bool: 如果是终端步骤返回True，否则返回False
    """
    assert ost_step is not None

    return "answer is" in ost_step.lower()


def print_tree_from_root(mcts_searcher, rollout_id, root_node, chosen_node=None, file=None):
    """
    从根节点开始打印MCTS树结构
    
    以树状结构可视化MCTS搜索树，显示每个节点的详细信息：
    - 节点类型和ID
    - MCTS统计信息（Q值、访问次数N、节点价值V、UCT值）
    - 节点具体内容（问题、答案、思考步骤等）
    - 终端解决方案节点会用红色高亮显示
    
    Args:
        mcts_searcher: MCTS搜索器实例，包含节点统计信息
        rollout_id (int): 当前rollout轮次ID
        root_node: 根节点
        chosen_node: 被选中的节点（会用额外标记显示）
        file: 输出文件对象，如果为None则输出到控制台
    """
    color_print = False if file else True

    def my_print(text):
        """内部打印函数，支持文件输出和控制台输出"""
        if file:
            file.write(text + "\n")
        else:
            print(text)

    def print_tree(parent_node, node, file, rollout_id):
        """递归打印树结构的内部函数"""
        to_print = ""

        num_indent = 4
        dash = "-" * num_indent * node.depth  # 根据深度生成缩进
        space = " " * num_indent * node.depth

        # 构建节点属性信息字符串
        attributes = f"Q: {round(mcts_searcher.Q[node], 2)}" + "; " + f"N: {mcts_searcher.N[node]}" + "; "
        attributes += f"V: {round(node.node_value, 2)}" if node.node_value is not None else "V: None"

        uct_value = "UCT: " + str(
            round(mcts_searcher._compute_uct(parent_node=parent_node, node=node, rollout_id=rollout_id), 2)
        )
        attributes += "; " + uct_value

        # 为终端解决方案节点添加标记
        solution_marker = "(T) " if node.is_valid_solution_node() else ""

        node_info = "[" + solution_marker + node.__str__() + ": " + attributes + "]"
        if chosen_node and node == chosen_node:
            node_info = "[" + node_info + "]"  # 被选中的节点添加双重括号
        node_info += " "

        # 为终端解决方案节点添加颜色高亮
        if color_print and node.is_valid_solution_node():
            node_details = Fore.RED + Style.BRIGHT + node_info + Fore.RESET + Style.RESET_ALL
        else:
            node_details = node_info

        # 根据节点类型添加具体内容
        if node.node_type is Node_Type.USER_QUESTION:
            gt = node.expected_answer.replace("\n", " ")
            node_details += f"User: {node.user_question}" + "\n" + space + " " * len(node_info) + f"Ground truth: {gt}"
        elif node.node_type is Node_Type.REPHRASED_USER_QUESTION:
            node_details += f"Reph-User: {node.user_question}"
        elif node.node_type is Node_Type.DIRECT_ANSWER:
            node_details += f"Ans: {node.direct_answer}"
        elif node.node_type is Node_Type.SUBQUESTION:
            node_details += f"Q: {node.subquestion}" + "\n" + space + " " * len(node_info) + f"A: {node.subanswer}"
        elif node.node_type is Node_Type.RE_SUBANSWER:
            node_details += f"Re-Ans: {node.re_subanswer}"
        elif node.node_type is Node_Type.OST_STEP:
            node_details += f"OST: {node.ost_step}"

        to_print += dash + node_details

        my_print(to_print)

        # 递归打印所有子节点
        for child in node.children:
            print_tree(node, child, file, rollout_id)

        # 根节点打印完成后添加分隔线
        if node.depth == 0:
            my_print("\n" + "=" * 50 + "\n")

    print_tree(parent_node=None, node=root_node, file=file, rollout_id=rollout_id)


def concat_subqs_and_subas(solution_trace: Dict[int, Dict[str, str]], question_index: int) -> Tuple[str, int]:
    """
    拼接子问题和子答案的轨迹字符串
    
    将解决方案轨迹中的所有子问题和子答案按格式拼接成字符串，
    用于生成新的子问题时提供上下文信息。
    
    Args:
        solution_trace: 解决方案轨迹字典，键为子问题ID，值为包含子问题和子答案的字典
        question_index: 问题索引，用于格式化输出
        
    Returns:
        tuple: (拼接后的轨迹字符串, 下一个子问题ID)
        
    示例输出格式：
        Question 1.1: 第一个子问题
        Answer 1.1: 第一个子答案
        Question 1.2: 第二个子问题
        Answer 1.2: 第二个子答案
    """
    solution_trace_str = ""

    for subquestion_id, solution_step in solution_trace.items():
        if subquestion_id == 0:  # 跳过根节点（用户问题）
            continue

        assert subquestion_id > 0
        assert "subquestion" in solution_step.keys() and "subanswer" in solution_step.keys()

        # 格式化子问题
        solution_trace_str += f"Question {question_index}." + str(subquestion_id) + ": " + solution_step["subquestion"]
        solution_trace_str += "\n"
        # 格式化子答案
        solution_trace_str += (
            f"Answer {question_index}." + str(subquestion_id) + ": " + solution_step["subanswer"]["text"]
        )
        solution_trace_str += "\n"

    # 计算下一个子问题的ID
    next_subquestion_id = int(sorted(solution_trace.keys())[-1]) + 1
    return solution_trace_str, next_subquestion_id


def concat_ost_steps(solution_trace: Dict[int, Dict[str, str]]) -> Tuple[str, int]:
    """
    拼接一步思考步骤的轨迹字符串
    
    将解决方案轨迹中的一步思考步骤按格式拼接成字符串，
    用于生成新的一步思考步骤时提供上下文信息。
    
    Args:
        solution_trace: 解决方案轨迹字典，键为步骤ID，值为包含步骤信息的字典
        
    Returns:
        tuple: (拼接后的一步思考步骤字符串, 下一个步骤ID)
        
    示例输出格式：
        Step 1: 第一个思考步骤
        Step 2: 第二个思考步骤
    """
    # 获取轨迹中的最后一个元组
    last_tuple = list(solution_trace.items())[-1]
    last_tuple_id, last_tuple_recording = last_tuple[0], last_tuple[1]
    
    # 确保最后一个记录包含一步思考步骤
    assert "ost_step" in last_tuple_recording.keys()
    
    # 如果存在一步思考步骤，则拼接所有步骤
    if len(last_tuple_recording["ost_step"]) > 0:
        solution_trace_str = ""
        for step_id, step_text in last_tuple_recording["ost_step"].items():
            solution_trace_str += f"Step {step_id}: " + step_text + "\n"
        return solution_trace_str, step_id + 1
    else:
        # 如果还没有一步思考步骤，返回空字符串和起始ID
        return "", 1


def concat_subqs_subas_as_ost_steps(solution_trace: Dict[int, Dict[str, str]]) -> Tuple[str, int]:
    """
    将子问题和子答案作为一步思考步骤拼接
    
    将解决方案轨迹中的子问题和子答案转换为一步思考步骤的格式，
    提取子答案中的核心推理内容（去掉"The answer is"部分）。
    
    Args:
        solution_trace: 解决方案轨迹字典
        
    Returns:
        tuple: (拼接后的步骤字符串, 下一个步骤ID)
        
    示例输入轨迹（子问题-子答案）：
    {
        "0": {
            "user_question": "Janet's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?",
            "ost_step": {}
        },
        "1": {
            "subquestion": " How many eggs do the ducks lay each day?",
            "subanswer": {
                "text": "The ducks lay 16 eggs per day. The answer is 16.",
                "value": 1.0
            },
            "ost_step": {}
        },
        "2": {
            "subquestion": " How many eggs does Janet eat or use for baking muffins?",
            "subanswer": {
                "text": "Janet eats 3 eggs for breakfast and uses 4 eggs for baking muffins. That's a total of 3 + 4 = 7 eggs. The answer is 7.",
                "value": 1.0
            },
            "ost_step": {}
        },
        "3": {
            "subquestion": " Now we can answer the question: How much in dollars does she make every day at the farmers' market?",
            "subanswer": {
                "text": "Since the ducks lay 16 eggs per day and Janet eats/use 7 eggs, she has 16 - 7 = 9 eggs left to sell at the market. Each egg is sold for $2, so she makes 9 * 2 = 18 dollars. The answer is 18.",
                "value": 1.0
            },
            "ost_step": {}
        }
    }

    期望输出：
        subqs_subas_as_ost_steps_str:

            Step 1: The ducks lay 16 eggs per day.
            Step 2: Janet eats 3 eggs for breakfast and uses 4 eggs for baking muffins. That's a total of 3 + 4 = 7 eggs.
            Step 3: Since the ducks lay 16 eggs per day and Janet eats/use 7 eggs, she has 16 - 7 = 9 eggs left to sell at the market. Each egg is sold for $2, so she makes 9 * 2 = 18 dollars.

        next_ost_step_id: 4
    """
    subqs_subas_as_ost_steps_str = ""
    step_id = 1
    
    # 遍历解决方案轨迹，提取子答案中的推理步骤
    while step_id in solution_trace:
        if "subanswer" in solution_trace[step_id]:
            # 使用正则表达式提取答案中的推理部分（去掉"The answer is"后的内容）
            match = re.search(r"(.+?\.) The answer is", solution_trace[step_id]["subanswer"]["text"])
            if match:
                # 如果找到匹配，提取推理部分
                step_text = match.group(1).strip()
            else:
                # 如果没有找到"The answer is"，使用完整文本
                step_text = solution_trace[step_id]["subanswer"]["text"].strip()
            
            # 格式化为步骤字符串
            subqs_subas_as_ost_steps_str += f"Step {step_id}: " + step_text + "\n"
            step_id += 1
        else:
            # 如果还没有子问题，返回空字符串和起始ID
            return "", 1
    
    return subqs_subas_as_ost_steps_str, step_id + 1


def concat_solution_trace(solution_trace: Dict[int, Dict[str, str]]):
    """
    拼接完整的解决方案轨迹
    
    将解决方案轨迹中的所有步骤拼接成完整的推理过程字符串，
    支持多种节点类型组合：子问题-子答案、一步思考、直接答案等。
    
    Args:
        solution_trace: 解决方案轨迹字典，包含完整的推理过程
        
    Returns:
        tuple: (完整轨迹字符串, 最终步骤字符串, 结束节点类型, 奖励值)
        
    注意：
        解决方案轨迹可能包含子问题-子答案和一步思考步骤的组合
        
    奖励值处理说明：
        - 使用 min(0, reward_value) + 1 确保奖励值至少为1
        - 原因：MCTS算法中奖励值为0会导致节点被认为是失败的，影响UCB计算
        - 通过+1确保即使是未完成的推理步骤也有基础奖励，维持探索平衡
    """
    solution_trace_str = ""
    final_step_str = ""
    end_node_type = None
    reward_value = 0.0

    # 遍历解决方案轨迹的每个项目
    for item_idx, (subq_id, solution_step) in enumerate(solution_trace.items()):
        if item_idx == 0:
            # 处理第一个项目（根节点）
            if len(solution_step["ost_step"]) == 0 and "direct_answer" in solution_step.keys():
                # 情况1：直接答案（无一步思考）
                # 
                # 输入例子：
                # solution_trace = {
                #     0: {
                #         "user_question": "What is 2 + 3?",
                #         "ost_step": {},
                #         "direct_answer": {
                #             "text": "2 + 3 = 5. The answer is 5.",
                #             "value": 1.0
                #         }
                #     }
                # }
                # 
                # 输出例子：
                # solution_trace_str = "Let's think step by step. 2 + 3 = 5. The answer is 5."
                # final_step_str = "2 + 3 = 5. The answer is 5."
                # end_node_type = Node_Type.DIRECT_ANSWER
                # reward_value = 1.0
                solution_trace_str += solution_step["direct_answer"]["text"].strip()
                final_step_str = solution_step["direct_answer"]["text"].strip()
                reward_value = (
                    solution_step["direct_answer"]["value"] if "value" in solution_step["direct_answer"] else 0.0
                )
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            elif len(solution_step["ost_step"]) > 0 and "direct_answer" in solution_step.keys():
                # 情况2：一步思考 + 直接答案
                # 
                # 输入例子：
                # solution_trace = {
                #     0: {
                #         "user_question": "What is 15 + 27?",
                #         "ost_step": {
                #             1: "I need to add 15 and 27.",
                #             2: "Let me break this down: 15 + 27 = 15 + 20 + 7 = 35 + 7 = 42."
                #         },
                #         "direct_answer": {
                #             "text": "The answer is 42.",
                #             "value": 1.0
                #         }
                #     }
                # }
                # 
                # 输出例子：
                # solution_trace_str = "Let's think step by step. I need to add 15 and 27. Let me break this down: 15 + 27 = 15 + 20 + 7 = 35 + 7 = 42. Now we can answer the question: The answer is 42."
                # final_step_str = "The answer is 42."
                # end_node_type = Node_Type.DIRECT_ANSWER
                # reward_value = 1.0
                for step_id, step_text in solution_step["ost_step"].items():
                    solution_trace_str += step_text.strip() + " "
                solution_trace_str += "Now we can answer the question: "
                solution_trace_str += solution_step["direct_answer"]["text"].strip()
                final_step_str = solution_step["direct_answer"]["text"].strip()
                reward_value = (
                    solution_step["direct_answer"]["value"] if "value" in solution_step["direct_answer"] else 0.0
                )
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            elif len(solution_step["ost_step"]) > 0 and "direct_answer" not in solution_step.keys():
                # 情况3：仅一步思考（无直接答案）
                # 
                # 输入例子：
                # solution_trace = {
                #     0: {
                #         "user_question": "How to solve this math problem?",
                #         "ost_step": {
                #             1: "First, I need to identify what the problem is asking.",
                #             2: "Then I should break it down into smaller steps.",
                #             3: "Finally, I can work through each step systematically."
                #         }
                #     }
                # }
                # 
                # 输出例子：
                # solution_trace_str = "Let's think step by step. First, I need to identify what the problem is asking. Then I should break it down into smaller steps. Finally, I can work through each step systematically."
                # final_step_str = "Finally, I can work through each step systematically."
                # end_node_type = Node_Type.OST_STEP
                # reward_value = 1.0 (min(0, 0.0) + 1)
                final_step_str = None
                for i, (step_id, step_text) in enumerate(solution_step["ost_step"].items()):
                    solution_trace_str += step_text.strip() + " "
                    if i == len(solution_step["ost_step"].items()) - 1:
                        final_step_str = step_text.strip()
                        reward_value = 0.0
                solution_trace_str = solution_trace_str.strip()
                end_node_type = Node_Type.OST_STEP
            else:
                # 如果都不满足，继续处理下一个项目
                continue
        elif 0 < item_idx < len(solution_trace) - 1:
            # 处理中间项目（子问题-子答案）
            # 提取子答案中的中间推理步骤（去掉"The answer is"部分）
            intermediate_step = solution_step["subanswer"]["text"].split("The answer is")[0].strip()
            solution_trace_str += intermediate_step + " "
            
            # 处理子问题后的一步思考步骤
            if len(solution_step["ost_step"]) > 0 and "direct_answer" in solution_step.keys():
                # 子问题 + 一步思考 + 直接答案
                for step_id, step_text in solution_step["ost_step"].items():
                    solution_trace_str += step_text.strip() + " "
                solution_trace_str += "Now we can answer the question: "
                solution_trace_str += solution_step["direct_answer"]["text"].strip()
                final_step_str = solution_step["direct_answer"]["text"].strip()
                reward_value = (
                    solution_step["direct_answer"]["value"] if "value" in solution_step["direct_answer"] else 0.0
                )
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            elif len(solution_step["ost_step"]) > 0 and "direct_answer" not in solution_step.keys():
                # 子问题 + 一步思考（无直接答案）
                final_step_str = None
                for i, (step_id, step_text) in enumerate(solution_step["ost_step"].items()):
                    solution_trace_str += step_text.strip() + " "
                    if i == len(solution_step["ost_step"].items()) - 1:
                        final_step_str = step_text.strip()
                        reward_value = 0.0
                solution_trace_str = solution_trace_str.strip()
                end_node_type = Node_Type.OST_STEP
        elif item_idx == len(solution_trace) - 1:
            # 处理最后一个项目（终端节点）
            assert item_idx > 0
            
            # 情况1：子问题-子答案（无一步思考，无直接答案）
            # 
            # 输入例子：
            # solution_trace = {
            #     0: {"user_question": "Janet's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?", "ost_step": {}},
            #     1: {"subquestion": "How many eggs do the ducks lay each day?", "subanswer": {"text": "The ducks lay 16 eggs per day. The answer is 16.", "value": 1.0}, "ost_step": {}},
            #     2: {"subquestion": "How many eggs does Janet eat or use for baking muffins?", "subanswer": {"text": "Janet eats 3 eggs for breakfast and uses 4 eggs for baking muffins. That's a total of 3 + 4 = 7 eggs. The answer is 7.", "value": 1.0}, "ost_step": {}},
            #     3: {"subquestion": "Now we can answer the question: How much in dollars does she make every day at the farmers' market?", "subanswer": {"text": "Since the ducks lay 16 eggs per day and Janet eats/use 7 eggs, she has 16 - 7 = 9 eggs left to sell at the market. Each egg is sold for $2, so she makes 9 * 2 = 18 dollars. The answer is 18.", "value": 1.0}, "ost_step": {}}
            # }
            # 
            # 输出例子：
            # solution_trace_str = "Let's think step by step. The ducks lay 16 eggs per day. Janet eats 3 eggs for breakfast and uses 4 eggs for baking muffins. That's a total of 3 + 4 = 7 eggs. Now we can answer the question: Since the ducks lay 16 eggs per day and Janet eats/use 7 eggs, she has 16 - 7 = 9 eggs left to sell at the market. Each egg is sold for $2, so she makes 9 * 2 = 18 dollars. The answer is 18."
            # final_step_str = "Since the ducks lay 16 eggs per day and Janet eats/use 7 eggs, she has 16 - 7 = 9 eggs left to sell at the market. Each egg is sold for $2, so she makes 9 * 2 = 18 dollars. The answer is 18."
            # end_node_type = Node_Type.SUBQUESTION
            # reward_value = 1.0
            if (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) == 0
                and "direct_answer" not in solution_step.keys()
            ):
                solution_trace_str += "Now we can answer the question: "
                solution_trace_str += solution_step["subanswer"]["text"].strip()
                final_step_str = solution_step["subanswer"]["text"].strip()
                reward_value = solution_step["subanswer"]["value"] if "value" in solution_step["subanswer"] else 0.0
                end_node_type = Node_Type.SUBQUESTION
                break
            # 情况2：子问题-子答案-一步思考（无直接答案）
            elif (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) > 0
                and "direct_answer" not in solution_step.keys()
            ):
                intermediate_step = solution_step["subanswer"]["text"].split("The answer is")[0].strip()
                solution_trace_str += intermediate_step + " "
                final_step_str = None
                for i, (step_id, step_text) in enumerate(solution_step["ost_step"].items()):
                    solution_trace_str += step_text.strip() + " "
                    if i == len(solution_step["ost_step"].items()) - 1:
                        final_step_str = step_text.strip()
                        reward_value = 0.0
                solution_trace_str = solution_trace_str.strip()
                end_node_type = Node_Type.OST_STEP
            # 情况3：子问题-子答案-一步思考-直接答案
            # 
            # 输入例子：
            # solution_trace = {
            #     0: {"user_question": "If John has 10 apples and gives away 3, how many does he have left?", "ost_step": {}},
            #     1: {"subquestion": "How many apples did John start with?", "subanswer": {"text": "John started with 10 apples. The answer is 10.", "value": 1.0}, "ost_step": {}},
            #     2: {"subquestion": "Now we can answer the question: How many apples does John have left?", "subanswer": {"text": "John had 10 apples and gave away 3. The answer is 7.", "value": 1.0}, "ost_step": {1: "Let me double-check this calculation: 10 - 3 = 7.", 2: "Yes, that's correct."}, "direct_answer": {"text": "John has 7 apples left. The answer is 7.", "value": 1.0}}
            # }
            # 
            # 输出例子：
            # solution_trace_str = "Let's think step by step. John started with 10 apples. Let me double-check this calculation: 10 - 3 = 7. Yes, that's correct. Now we can answer the question: John has 7 apples left. The answer is 7."
            # final_step_str = "John has 7 apples left. The answer is 7."
            # end_node_type = Node_Type.DIRECT_ANSWER
            # reward_value = 1.0
            elif (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) > 0
                and "direct_answer" in solution_step.keys()
            ):
                intermediate_step = solution_step["subanswer"]["text"].split("The answer is")[0].strip()
                solution_trace_str += intermediate_step + " "
                for step_id, step_text in solution_step["ost_step"].items():
                    solution_trace_str += step_text.strip() + " "
                solution_trace_str += "Now we can answer the question: "
                solution_trace_str += solution_step["direct_answer"]["text"].strip()
                final_step_str = solution_step["direct_answer"]["text"].strip()
                reward_value = (
                    solution_step["direct_answer"]["value"] if "value" in solution_step["direct_answer"] else 0.0
                )
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            # 情况4：子问题-子答案-直接答案（无一步思考）
            elif (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) == 0
                and "direct_answer" in solution_step.keys()
            ):
                intermediate_step = solution_step["subanswer"]["text"].split("The answer is")[0].strip()
                solution_trace_str += intermediate_step + " "
                solution_trace_str += "Now we can answer the question: "
                solution_trace_str += solution_step["direct_answer"]["text"].strip()
                final_step_str = solution_step["direct_answer"]["text"].strip()
                reward_value = (
                    solution_step["direct_answer"]["value"] if "value" in solution_step["direct_answer"] else 0.0
                )
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            # 情况5：仅直接答案
            elif "direct_answer" in solution_step.keys():
                assert len(solution_step["ost_step"]) == 0
                assert "subanswer" not in solution_step.keys()

                solution_trace_str += "Now we can answer the question: "
                solution_trace_str += solution_step["direct_answer"]["text"].strip()
                final_step_str = solution_step["direct_answer"]["text"].strip()
                reward_value = (
                    solution_step["direct_answer"]["value"] if "value" in solution_step["direct_answer"] else 0.0
                )
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            else:
                # 如果没有匹配的情况，进入调试模式
                import pdb
                pdb.set_trace()

    # 标准化轨迹字符串格式
    solution_trace_str = solution_trace_str.replace("Let's think step by step. ", "")
    solution_trace_str = "Let's think step by step. " + solution_trace_str

    # 奖励值处理：确保奖励值至少为1
    # 原因：
    # 1. MCTS算法中，奖励值为0会被视为失败节点，影响UCB值计算
    # 2. 在UCB公式中：UCB = Q + C * sqrt(ln(N_parent) / N_child)
    #    其中Q是平均奖励，如果Q=0，节点的吸引力会大幅下降
    # 3. 通过min(0, reward_value) + 1确保：
    #    - 负奖励值变为1（基础奖励）：min(0, -0.5) + 1 = -0.5 + 1 = 0.5
    #    - 0奖励值变为1（基础奖励）：min(0, 0.0) + 1 = 0.0 + 1 = 1.0
    #    - 正奖励值保持不变：min(0, 1.5) + 1 = 0.0 + 1 = 1.0
    # 4. 实际上这个公式的效果是：
    #    - 如果reward_value <= 0，则奖励值 = reward_value + 1
    #    - 如果reward_value > 0，则奖励值 = 1.0
    # 5. 这样设计的目的是给所有轨迹一个基础奖励值1，避免0奖励影响MCTS探索
    # 6. 对于OST_STEP等中间步骤，由于reward_value=0.0，最终得到1.0的基础奖励
    # 7. 对于有具体评估分数的终端节点，统一标准化为1.0，简化奖励计算
    return solution_trace_str.strip(), final_step_str.strip(), end_node_type, min(0, reward_value) + 1


def concat_rap_solution_trace(solution_trace: str):
    """
    拼接RAP（Reasoning via Planning）解决方案轨迹
    
    从RAP格式的解决方案轨迹中提取答案列表，
    去掉格式化标记和最终答案数字。
    
    Args:
        solution_trace: RAP格式的解决方案轨迹字符串
        
    Returns:
        tuple: (拼接后的答案字符串, 最终步骤字符串)
    """
    solution_trace_list = solution_trace.split("\n")
    answer_list = []
    
    # 遍历轨迹列表，提取答案行
    for item in solution_trace_list:
        if item.startswith("Answer"):
            # 去掉答案编号格式（如"Answer 1.1: "）
            item = re.sub(r"Answer \d+\.\d+: ", "", item)
            final_step = item
            # 去掉最终答案数字（如" The answer is 16."）
            item = re.sub(r" The answer is \d+\.", "", item)
            answer_list.append(item)
    
    return " ".join(answer_list).strip(), final_step


def concat_subq_suba_trace(solution_trace: Dict[int, Dict[str, str]]):
    """
    拼接子问题-子答案轨迹
    
    将解决方案轨迹转换为结构化的子问题-子答案格式，
    用于特定的推理模式。
    
    Args:
        solution_trace: 解决方案轨迹字典
        
    Returns:
        tuple: (结构化轨迹列表, 最终步骤字符串, 结束节点类型)
        
    注意：
        解决方案轨迹可能包含子问题-子答案和一步思考步骤
    """
    solution_trace_str = ["Let's think step by step."]
    final_step_str = ""
    end_node_type = None

    # 遍历解决方案轨迹
    for item_idx, (subq_id, solution_step) in enumerate(solution_trace.items()):
        if item_idx == 0:
            # 第一个项目：确保格式正确
            assert len(solution_step["ost_step"]) == 0
            assert "direct_answer" not in solution_step.keys()
        elif 0 < item_idx < len(solution_trace) - 1:
            # 中间项目：添加子问题-子答案对
            assert len(solution_step["ost_step"]) == 0
            solution_trace_str.append(
                {"subq": solution_step["subquestion"], "suba": solution_step["subanswer"]["text"]}
            )
            # # concat trace for one-step thought step after subquestion
            # if len(solution_step["ost_step"]) > 0 and "direct_answer" in solution_step.keys():
            #     for step_id, step_text in solution_step["ost_step"].items():
            #         solution_trace_str += step_text.strip() + " "
            #     solution_trace_str += "Now we can answer the question: "
            #     solution_trace_str += solution_step["direct_answer"]["text"].strip()
            #     final_step_str = solution_step["direct_answer"]["text"].strip()
            #     end_node_type = Node_Type.DIRECT_ANSWER
            #     break
            # elif len(solution_step["ost_step"]) > 0 and "direct_answer" not in solution_step.keys():
            #     final_step_str = None
            #     for i, (step_id, step_text) in enumerate(solution_step["ost_step"].items()):
            #         solution_trace_str += step_text.strip() + " "
            #         if i == len(solution_step["ost_step"].items()) - 1:
            #             final_step_str = step_text.strip()
            #     solution_trace_str = solution_trace_str.strip()
            #     end_node_type = Node_Type.OST_STEP
        elif item_idx == len(solution_trace) - 1:
            # 最后一个项目：处理不同的结束情况
            assert item_idx > 0
            
            # 情况1：子问题-子答案（标准结束）
            # 1. subq-suba
            if (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) == 0
                and "direct_answer" not in solution_step.keys()
            ):
                solution_trace_str.append(
                    {"subq": solution_step["subquestion"], "suba": solution_step["subanswer"]["text"]}
                )
                final_step_str = solution_step["subanswer"]["text"].strip()
                end_node_type = Node_Type.SUBQUESTION
                break
            # 情况2和3：包含一步思考的复杂情况（当前不支持）
            # 2. subq-suba-ost
            elif (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) > 0
                and "direct_answer" not in solution_step.keys()
            ):
                assert False
            # 3. subq-suba-ost-diranswer
            elif (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) > 0
                and "direct_answer" in solution_step.keys()
            ):
                assert False
            # 情况4：子问题-子答案-直接答案
            # 4. subq-suba-diranswer
            elif (
                "subanswer" in solution_step.keys()
                and len(solution_step["ost_step"]) == 0
                and "direct_answer" in solution_step.keys()
            ):
                solution_trace_str.append(
                    {"subq": solution_step["subquestion"], "suba": solution_step["subanswer"]["text"]}
                )
                solution_trace_str.append(solution_step["direct_answer"]["text"])
                final_step_str = solution_step["direct_answer"]["text"].strip()
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            # 情况5：仅直接答案
            # 5. diranswer
            elif "direct_answer" in solution_step.keys():
                assert len(solution_step["ost_step"]) == 0
                assert "subanswer" not in solution_step.keys()

                solution_trace_str.append(
                    solution_step["direct_answer"]["text"].replace("Let's think step by step.", "")
                )
                final_step_str = solution_step["direct_answer"]["text"].strip()
                end_node_type = Node_Type.DIRECT_ANSWER
                break
            else:
                # 未匹配的情况，进入调试模式
                import pdb
                pdb.set_trace()

    return solution_trace_str, final_step_str.strip(), end_node_type


def mask_solution_trace(
    solution_trace_str: str, num_return: int, left_boundary: float, right_boundary: float
) -> list[str]:
    """
    掩码解决方案轨迹
    
    根据指定的边界和返回数量，将完整的解决方案轨迹按比例截断，
    生成多个不同长度的部分轨迹。
    
    Args:
        solution_trace_str: 完整的解决方案轨迹字符串
        num_return: 需要返回的掩码轨迹数量
        left_boundary: 左边界比例（0.0-1.0）
        right_boundary: 右边界比例（0.0-1.0）
        
    Returns:
        list[str]: 不同长度的掩码轨迹列表
        
    示例：
        输入: "opasdjifpoaisdfjpoasidfjapsodifj", num_return=4, left=0.2, right=0.8
        输出: ["opasd", "opasdjifp", "opasdjifpoaisdfj", "opasdjifpoaisdfjpoasidfjaps"]
    """
    # opasdjifpoaisdfjpoasidfjapsodifj, num_return: 4, left: 0.2, right: 0.8
    # return: opasd, opasdjifp, opasdjifpoaisdfj, opasdjifpoaisdfjpoasidfjaps
    if num_return == 1:
        # 如果只需要一个结果，间隔为0
        interval = 0
    else:
        assert num_return > 1
        assert right_boundary >= left_boundary, f"right_boundary: {right_boundary} < left_boundary: {left_boundary}"
        # 计算边界间隔
        interval = (right_boundary - left_boundary) / (num_return - 1)

    # 将轨迹字符串按空格分割为单词列表
    words_in_solution_trace = solution_trace_str.split(" ")
    ost_len = len(words_in_solution_trace)
    
    # 从最少到最多生成掩码轨迹
    # Mask the solution trace string from least to most
    masked_solution_traces = []
    for i in range(num_return):
        # 计算当前前缀部分的比例
        prefix_part_ratio = left_boundary + i * interval
        # 计算前缀部分的单词数量（向上取整）
        prefix_part_num_words = math.ceil(ost_len * prefix_part_ratio)
        # 生成前缀部分字符串
        prefix_part_str = " ".join(words_in_solution_trace[:prefix_part_num_words])
        masked_solution_traces.append(prefix_part_str)

    return masked_solution_traces


def mask_subq_suba_trace(solution_trace_str: list, num_return: int, evaluator: Evaluator) -> list[str]:
    """
    掩码子问题-子答案轨迹
    
    从结构化的子问题-子答案轨迹中生成部分轨迹，
    用于渐进式推理和答案提取。
    
    Args:
        solution_trace_str: 结构化的解决方案轨迹列表
        num_return: 需要返回的掩码轨迹数量
        evaluator: 评估器实例，用于提取答案
        
    Returns:
        list[str]: 掩码轨迹列表，每个元素包含[前缀字符串, 当前答案]
        
    注意：
        当前实现主要支持num_return=1的情况，多返回值功能待实现
    """
    # opasdjifpoaisdfjpoasidfjapsodifj, num_return: 4, left: 0.2, right: 0.8
    # return: opasd, opasdjifp, opasdjifpoaisdfj, opasdjifpoaisdfjpoasidfjaps
    if num_return == 1:
        interval = 0
    else:
        assert num_return > 1
    # TODO: 实现num_return > 1的情况

    # 从最少到最多生成掩码轨迹
    # Mask the solution trace string from least to most
    masked_solution_traces = []
    for i in range(1, len(solution_trace_str)):
        if "subq" in solution_trace_str[i]:
            # 处理子问题-子答案对
            prefix_part_str = (
                "\n".join(
                    [
                        f"Subquestion: {item['subq']}\nSubanswer: {item['suba']}" if "subq" in item else item
                        for item in solution_trace_str[:i]
                    ]
                )
                + "\nSubquestion: "
                + solution_trace_str[i]["subq"]
                + "\nSubanswer: "
            )
            # 提取当前答案
            curr_answer = evaluator.extract_answer_from_model_completion(solution_trace_str[i]["suba"])
            masked_solution_traces.append([prefix_part_str, curr_answer])
        else:
            # 处理最终答案
            prefix_part_str = (
                "\n".join(
                    [
                        f"Subquestion: {item['subq']}\nSubanswer: {item['suba']}" if "subq" in item else item
                        for item in solution_trace_str[:i]
                    ]
                )
                + "\nSubquestion: Now we can answer the question:"
            )
            # 提取当前答案
            curr_answer = evaluator.extract_answer_from_model_completion(solution_trace_str[i])
            masked_solution_traces.append([prefix_part_str, curr_answer])
    
    return masked_solution_traces


def make_hint(
    solution_trace: Dict[int, Dict[str, str]], node_type: Node_Type, new_subq=None, new_suba=None, new_ost_step=None
) -> str:
    """
    生成提示信息
    
    根据当前解决方案轨迹和节点类型，生成用于指导后续推理的提示信息。
    
    Args:
        solution_trace: 当前解决方案轨迹
        node_type: 节点类型，决定提示格式
        new_subq: 新的子问题（可选）
        new_suba: 新的子答案（可选）
        new_ost_step: 新的一步思考步骤（可选）
        
    Returns:
        str: 格式化的提示字符串
        
    Raises:
        ValueError: 当节点类型无效时抛出异常
    """
    if node_type in [Node_Type.SUBQUESTION, Node_Type.RE_SUBANSWER]:
        # 处理子问题和重新子答案类型
        hint = ""

        # 遍历解决方案轨迹，生成提示
        for subquestion_id, solution_step in solution_trace.items():
            if subquestion_id == 0:
                # 跳过根节点
                continue

            assert subquestion_id > 0
            assert "subquestion" in solution_step.keys() and "subanswer" in solution_step.keys()

            # 格式化提示：Hint {id}: {子问题} {子答案}
            hint += f"Hint " + str(subquestion_id) + ": " + solution_step["subquestion"]
            hint += " "
            hint += solution_step["subanswer"]["text"]
            hint += "\n"

        # 如果有新的子问题和子答案，添加到提示中
        if new_subq is not None and new_suba is not None:
            hint += f"Hint {len(solution_trace)}: " + new_subq + " " + new_suba

        hint = hint.strip("\n")
    elif node_type is Node_Type.OST_STEP:
        # 处理一步思考类型
        hint = "Hint: "
        last_tuple = list(solution_trace.items())[-1]
        last_tuple_recording = last_tuple[1]
        assert last_tuple_recording["ost_step"]
        
        # 拼接所有一步思考步骤
        for step_id, step_text in last_tuple_recording["ost_step"].items():
            hint += step_text + " "

        # 如果有新的一步思考步骤，添加到提示中
        if new_ost_step is not None:
            hint += new_ost_step

        hint = hint.strip(" ")
    else:
        raise ValueError(f"Invalid node type: {node_type}.")

    return hint


def make_response_prefix(
    solution_trace: Dict[int, Dict[str, str]], node_type: Node_Type, new_subq=None, new_suba=None, new_ost_step=None
) -> str:
    """
    生成响应前缀
    
    根据当前解决方案轨迹和节点类型，生成用于模型响应的前缀字符串。
    
    Args:
        solution_trace: 当前解决方案轨迹
        node_type: 节点类型，决定前缀格式
        new_subq: 新的子问题（可选）
        new_suba: 新的子答案（可选）
        new_ost_step: 新的一步思考步骤（可选）
        
    Returns:
        str: 格式化的响应前缀字符串
        
    Raises:
        ValueError: 当节点类型无效时抛出异常
    """
    if node_type in [Node_Type.SUBQUESTION, Node_Type.RE_SUBANSWER]:
        # 处理子问题和重新子答案类型
        response_prefix = ""
        answer_marker = "The answer is"  # todo: hard code "The answer is"

        # 遍历解决方案轨迹，提取推理部分（去掉答案标记）
        for subquestion_id, solution_step in solution_trace.items():
            if subquestion_id == 0:
                # 跳过根节点
                continue

            assert subquestion_id > 0
            assert "subquestion" in solution_step.keys() and "subanswer" in solution_step.keys()

            # 提取子答案中的推理部分（去掉"The answer is"后的内容）
            response_prefix += solution_step["subanswer"]["text"].split(answer_marker)[0]
            response_prefix += " "

        # 如果有新的子问题和子答案，处理新的子答案
        if new_subq is not None and new_suba is not None:
            response_prefix += new_suba.split(answer_marker)[0]

        response_prefix = response_prefix.strip(" ")
    elif node_type is Node_Type.OST_STEP:
        # 处理一步思考类型
        response_prefix = ""

        # 获取最后一个轨迹记录
        last_tuple = list(solution_trace.items())[-1]
        last_tuple_recording = last_tuple[1]
        
        # 如果存在一步思考步骤，拼接所有步骤
        if "ost_step" in last_tuple_recording.keys():
            for step_id, step_text in last_tuple_recording["ost_step"].items():
                response_prefix += step_text + " "

        # 如果有新的一步思考步骤，添加到前缀中
        if new_ost_step is not None:
            response_prefix += new_ost_step

        response_prefix = response_prefix.strip(" ")
    elif node_type is None and solution_trace is None:
        # 空情况处理
        response_prefix = ""
    else:
        raise ValueError(f"Invalid node type: {node_type}.")

    # 添加标准的思考前缀
    think = "Let's think step by step. "
    return think + response_prefix if think not in response_prefix else response_prefix


def find_valid_solution_nodes(root_node):
    """
    查找有效的解决方案节点
    
    从MCTS树中递归查找所有有效的解决方案节点。
    有效的解决方案节点是指已经得出最终答案的终端节点。
    
    Args:
        root_node: MCTS树的根节点
        
    Returns:
        list: 包含所有有效解决方案节点的列表
    """
    valid_solution_nodes = []

    def recursion(node):
        """
        递归遍历节点的内部函数
        
        Args:
            node: 当前遍历的节点
        """
        if node.is_valid_solution_node():
            # 如果是有效解决方案节点，添加到列表中
            valid_solution_nodes.append(node)
            return

        if not node.children:  #! no children
            # 如果没有子节点，返回
            return

        # 递归遍历所有子节点
        for child in node.children:
            recursion(child)

    # 从根节点开始递归遍历
    recursion(root_node)

    return valid_solution_nodes


def find_best_solution(root_node, evaluator, enable_potential_score=False):
    """
    从MCTS树中找到最佳解决方案
    
    【核心选择机制】：
    从所有有效的解决方案节点中选择最佳答案，支持两种选择策略：
    1. 基于置信度的选择（默认）：纯粹基于答案出现频次
    2. 基于潜在分数的选择：考虑推理路径的一致性
    
    【潜在分数评估原理】：
    潜在分数基于"推理路径一致性"的假设：
    - 如果一个解决方案的推理路径在各个深度都能产生与最终答案一致的潜在答案
    - 那么这个解决方案更可能是正确的
    - 计算方式：将各深度的一致性比例相乘得到潜在分数
    
    【算法流程】：
    1. 查找所有有效解决方案节点
    2. 提取每个节点的解决方案内容
    3. 如果启用潜在分数：计算每个节点的潜在分数作为权重
    4. 使用评估器的投票机制选择最佳答案
    
    【判断依据】：
    - 无潜在分数：答案频次（多数投票）
    - 有潜在分数：加权频次（潜在分数 × 频次）
    
    Args:
        root_node: MCTS树的根节点
        evaluator: 评估器实例，用于提取和比较答案
        enable_potential_score: 是否启用潜在分数评估
        
    Returns:
        tuple: (最佳答案, 最佳解决方案, 置信度, 对应节点, 所有解决方案节点)
               如果没有找到解决方案则返回 (None, None)
               
    示例：
        假设有3个解决方案节点：
        - 节点A：答案="42"，潜在分数=0.8
        - 节点B：答案="42"，潜在分数=0.6  
        - 节点C：答案="43"，潜在分数=0.9
        
        不启用潜在分数：选择"42"（频次2/3=0.67 > 1/3=0.33）
        启用潜在分数：计算加权分数后选择最高分答案
    """
    # 查找所有有效的解决方案节点
    solution_nodes = find_valid_solution_nodes(root_node)

    if len(solution_nodes) == 0:
        # 如果没有找到解决方案节点，返回None
        return None, None

    def extract_solution_from_node(node):
        """
        从节点中提取解决方案内容
        
        根据节点类型提取相应的解决方案内容：
        - SUBQUESTION节点：提取子答案
        - DIRECT_ANSWER节点：提取直接答案
        - 其他类型：返回None
        
        Args:
            node: 解决方案节点
            
        Returns:
            str: 解决方案内容
        """
        if node.node_type is Node_Type.SUBQUESTION:
            return node.subanswer
        elif node.node_type is Node_Type.DIRECT_ANSWER:
            return node.direct_answer
        else:
            return None

    # 从所有解决方案节点中提取解决方案
    solutions = [extract_solution_from_node(node) for node in solution_nodes]

    def calculate_potential_score_for_solution_node(node):
        """
        计算解决方案节点的潜在分数
        
        【潜在分数计算原理】：
        潜在分数基于推理路径中各深度的潜在答案与最终答案的一致性：
        1. 遍历节点的潜在答案历史（每个深度的潜在答案列表）
        2. 计算每个深度的潜在答案与最终答案的匹配度
        3. 将所有深度的匹配度相乘得到潜在分数
        
        【计算公式】：
        潜在分数 = ∏(深度i的一致性比例) for i in 所有深度
        深度i的一致性比例 = 与最终答案一致的潜在答案数 / 该深度的潜在答案总数
        
        【示例】：
        假设推理路径有3个深度：
        - 深度1：潜在答案["40", "42", "44"]，与最终答案"42"一致的有1个 → 1/3
        - 深度2：潜在答案["41", "42"]，与最终答案"42"一致的有1个 → 1/2  
        - 深度3：最终答案"42"，跳过
        潜在分数 = (1/3) × (1/2) = 1/6 = 0.167
        
        Args:
            node: 解决方案节点
            
        Returns:
            float: 潜在分数（0.0-1.0）
        """
        # 提取模型的最终答案
        model_answer = evaluator.extract_answer_from_model_completion(extract_solution_from_node(node))
        # 获取潜在答案历史记录 {深度 -> [潜在答案列表]}
        potential_answers_history = node.potential_answers_history  # {depth -> [potential answers]}
        # 确保当前深度的潜在答案为None（已经是最终答案）
        assert potential_answers_history[node.depth] is None

        potential_score = 1
        # 遍历所有深度的潜在答案
        for depth, depth_potential_answers in potential_answers_history.items():
            if depth < node.depth:
                # 计算当前深度潜在答案与最终答案的匹配度
                depth_score = sum(
                    evaluator.check_answers_equiv(dpa, model_answer) for dpa in depth_potential_answers
                ) / len(depth_potential_answers)
                # 将匹配度累乘到潜在分数中
                potential_score *= depth_score

        # 设置节点的潜在分数
        node.set_potential_score(potential_score)
        return potential_score

    # 根据是否启用潜在分数来计算先验权重
    prior_weights = (
        [calculate_potential_score_for_solution_node(node) for node in solution_nodes]
        if enable_potential_score
        else None
    )
    
    # 使用评估器找到最有信心的答案
    # 如果有先验权重，会使用加权投票；否则使用纯频次投票
    top_answer, top_completion, top_completion_id, top_confidence = evaluator.find_most_confident_answer(
        solutions, prior_weights
    )
    return top_answer, top_completion, top_confidence, solution_nodes[top_completion_id], solution_nodes


def stochastic_find_best_solution(
    root_node,
    evaluator,
    enable_potential_score,
):
    # todo: what strategy do we use to select best node?
    """
    随机选择最佳解决方案
    
    使用随机策略从MCTS树中选择最佳解决方案，支持潜在分数权重。
    与确定性选择不同，该函数引入随机性来选择解决方案。
    
    该函数从MCTS树中的解决方案节点中找到最佳解决方案。
    
    Args:
        root_node: MCTS树的根节点
        evaluator: 评估器实例，用于提取和比较答案
        enable_potential_score: 是否启用潜在分数评估
        
    Returns:
        tuple: (最佳答案, 最佳解决方案, 置信度, 对应节点, 所有解决方案节点, 所有解决方案)
               如果没有找到解决方案则返回 (None, None)
               
    注意：
        该函数的选择策略是什么？需要根据具体实现来确定。
    """
    # 查找所有有效的解决方案节点
    solution_nodes = find_valid_solution_nodes(root_node)

    if len(solution_nodes) == 0:
        # 如果没有找到解决方案节点，返回None
        return None, None

    def extract_solution_from_node(node):
        """
        从节点中提取解决方案内容
        
        Args:
            node: 解决方案节点
            
        Returns:
            解决方案内容（子答案或直接答案）
        """
        if node.node_type is Node_Type.SUBQUESTION:
            return node.subanswer
        elif node.node_type is Node_Type.DIRECT_ANSWER:
            return node.direct_answer
        else:
            return None

    # 从所有解决方案节点中提取解决方案
    solutions = [extract_solution_from_node(node) for node in solution_nodes]

    def calculate_potential_score_for_solution_node(node):
        """
        计算解决方案节点的潜在分数
        
        计算方法与find_best_solution中的相同，基于推理路径一致性。
        
        Args:
            node: 解决方案节点
            
        Returns:
            float: 潜在分数（0.0-1.0）
        """
        # 提取模型的最终答案
        model_answer = evaluator.extract_answer_from_model_completion(extract_solution_from_node(node))
        # 获取潜在答案历史记录 {深度 -> [潜在答案列表]}
        potential_answers_history = node.potential_answers_history  # {depth -> [potential answers]}
        # 确保当前深度的潜在答案为None（已经是最终答案）
        assert potential_answers_history[node.depth] is None

        potential_score = 1
        # 遍历所有深度的潜在答案
        for depth, depth_potential_answers in potential_answers_history.items():
            if depth < node.depth:
                # 计算当前深度潜在答案与最终答案的匹配度
                depth_score = sum(
                    evaluator.check_answers_equiv(dpa, model_answer) for dpa in depth_potential_answers
                ) / len(depth_potential_answers)
                # 将匹配度累乘到潜在分数中
                potential_score *= depth_score

        # 设置节点的潜在分数
        node.set_potential_score(potential_score)
        return potential_score

    # 根据是否启用潜在分数来计算先验权重
    prior_weights = (
        [calculate_potential_score_for_solution_node(node) for node in solution_nodes]
        if enable_potential_score
        else None
    )
    
    # 使用评估器的随机选择方法找到最有信心的答案
    top_answer, top_completion, top_completion_id, top_confidence = evaluator.stochastic_find_most_confident_answer(
        completions=solutions, prior_weights=prior_weights
    )
    return top_answer, top_completion, top_confidence, solution_nodes[top_completion_id], solution_nodes, solutions
